import java.util.Arrays;
import java.util.PriorityQueue;

/**
 * 面试题40：最小的k个数
 */
public class Offer_40 {
    /**
     * 方法四：快排思想
     * <p>
     * 时间复杂度：O(n)
     * <p>
     * 空间复杂度：O(logn)
     */
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k >= arr.length) {
            return arr;
        }
        return quickSearch(arr, k, 0, arr.length - 1);
    }

    private int[] quickSearch(int[] arr, int k, int left, int right) {
        int i = left, j = right;
        while (i < j) {
            while (i < j && arr[j] >= arr[left]) {
                j--;
            }
            while (i < j && arr[i] <= arr[left]) {
                i++;
            }
            swap(arr, i, j);
        }
        swap(arr, i, left);
        if (i > k) {
            return quickSearch(arr, k, left, i - 1);
        }
        if (i < k) {
            return quickSearch(arr, k, i + 1, right);
        }
        return Arrays.copyOf(arr, k);
    }

    private void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    /**
     * 方法三：大根堆
     * <p>
     * 时间复杂度：O(nlogk)
     * <p>
     * 空间复杂度：O(k)
     */
    public int[] getLeastNumbers3(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }

        // 大根堆，初始容量为 k
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(k, (a, b) -> (b - a));
        for (int i = 0; i < k; i++) {
            priorityQueue.offer(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            if (arr[i] < priorityQueue.peek()) {
                priorityQueue.poll();
                priorityQueue.offer(arr[i]);
            }
        }

        int[] ans = new int[k];
        for (int i = 0; i < k; i++) {
            ans[i] = priorityQueue.poll();
        }
        return ans;
    }

    /**
     * 方法二：统计每个数字出现的次数（适用于数字范围较小）
     * <p>
     * 时间复杂度：O(n)
     * <p>
     * 空间复杂度：取决于数据范围
     */
    public int[] getLeastNumbers2(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }

        // 统计每个数字出现的次数
        // 0 <= arr[i] <= 10000
        int[] count = new int[10001];
        for (int num : arr) {
            count[num]++;
        }

        // 根据 count 数组从头找出 k 个数作为返回结果
        int[] ans = new int[k];
        int index = 0;
        for (int i = 0; i < count.length; i++) {
            while (count[i]-- > 0 && index < k) {
                ans[index++] = i;
            }
            if (index == k) {
                break;
            }
        }
        return ans;
    }

    /**
     * 方法一：排序
     * <p>
     * 时间复杂度：O(nlogn)
     * <p>
     * 空间复杂度：O(logn)
     */
    public int[] getLeastNumbers1(int[] arr, int k) {
        Arrays.sort(arr);
        int[] ans = new int[k];
        for (int i = 0; i < k; i++) {
            ans[i] = arr[i];
        }
        return ans;
    }
}
